import pandas as pd 
import os
import numpy as np 
import warnings
warnings.filterwarnings('ignore')

# 策略逻辑 截面开仓
# 阈值以上计算排名，买入n只，开仓至金额不够

# 买入仓位表， 当日开盘价买入
# 卖出卖出表， 当日收盘价卖出
# 持仓表

# 到期平仓判断
# 开仓判断
# 止盈止损

#阈值加rank开仓方式，附带止盈止损
class Pos_make_rank_thred_step_combine():   
    
    def __init__(self, args, open_root=None, close_root=None):
        self.args = args
        self.money_init = args.money_init
        self.money = args.money_init
        self.money = args.money_init
        self.open_thred = args.open_thred
        self.vol_num = 20
        self.predict = self.get_predict()   
        self.open, self.close = self.get_price(open_root, close_root)
        self.hold_pos = self.pos_init()
        self.buy_pos = self.pos_init()
        self.sell_pos = self.pos_init()
        self.cost_df= self.pos_init()
        self.money_df = pd.DataFrame(index=self.predict.index, columns=['money'])
        self.stop_ret = args.stop_ret
        self.stop_loss = args.stop_loss
        self.cost_ratio = args.cost_ratio
        self.hold_num = args.hold_num
        self.pos_root = self.pos_root_make()

    
    
    def pos_root_make(self):
        result_root = os.path.join(self.args.result_root, self.__class__.__name__, f'loss{self.stop_loss}_ret{self.stop_ret}_thred{self.open_thred}_holdnum{self.hold_num}')
        if not os.path.exists(result_root):
            os.makedirs(result_root)
        return result_root
    
    def log_write(self, content):
        with open(os.path.join(self.pos_root, 'log.txt'),"a") as file:   #只需要将之前的”w"改为“a"即可，代表追加内容
            file.write(content)

    def get_predict(self):
        combine_lst = [5,10,30]
        df = 0
        vol = 0
        for t in combine_lst:
            df_i = self.get_predict_t(t)
            # df = df + df_i
            df = df + df_i / (df_i.rolling(10, min_periods=1).std() +0.000001)
            # vol = vol + (1/ df_i.rolling(5, min_periods=1).std())
        df = df.shift(1)
        df = df / len(combine_lst)
        # vol = vol.shift(1)
        df = df / vol
        return df
    
    def get_price(self, open_root, close_root):  
        # sourcery skip: avoid-builtin-shadow
        if open_root is None:
            open_root = r'data\stock_data\consentrate_daily_price\open.pkl.gzip'
        if close_root is None:
            close_root = r'data\stock_data\consentrate_daily_price\close.pkl.gzip'
        open = pd.read_pickle(open_root)
        close = pd.read_pickle(close_root)
        open = open.loc[self.predict.index, self.predict.columns]
        close = close.loc[self.predict.index, self.predict.columns]
        open = open.fillna(method='ffill')
        close = close.fillna(method='ffill')
        return open, close

    def get_predict_t(self, t):
        read_root = self.args.result_root[:-2]+str(t)
        df = pd.read_pickle(os.path.join(read_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df
    
    def pos_init(self):
        contract = self.predict.columns
        zero_ary = np.zeros((len(self.predict),len(contract)))
        pos_df = pd.DataFrame(zero_ary, index= self.predict.index, columns=contract)
        pos_df = pos_df.astype(int)
        return pos_df
    
    def pos_make(self):         
        for idx, idv in enumerate(self.predict.index):
            # 到期平仓判断
            # self.time_stop_pos(idx, idv)

            # 开仓判断
            print(idv)
            self.log_write(f'\n{idv}')
            self.pos_add(idx, idv)
            self.stop_ratio(idx, idv)
            self.money_df.loc[idv] = self.money
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        self.pos_result_save()
    
    def pos_result_save(self):
        self.buy_pos.to_csv(f'{self.pos_root}/buy_pos.csv')
        self.sell_pos.to_csv(f'{self.pos_root}/sell_pos.csv')
        self.hold_pos.to_csv(f'{self.pos_root}/hold_pos.csv')
        self.money_df.to_csv(f'{self.pos_root}/money.csv')

    def time_stop_pos(self, idx, idv):
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        for code in self.columns:
            hold_pos = self.hold_pos.loc[idv, code].values
            if  hold_pos > 0:
                self.sell_pos.loc[idv, code] = hold_pos

    def pos_add(self, idx, idv):
        pred = self.predict.loc[idv].sort_values(ascending=False)
        pred = pred[:self.hold_num]
        pred = pred[pred>=self.open_thred]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        open_command = any(pred[code] >= self.open_thred for code in buy_codes)
        if open_command:
            while (self.money > cost_lst+price_list).any():
                for code in buy_codes:#debug 检查此处
                    if pred[code] >= self.open_thred: #满足开仓信号要求                       
                        price_amount = self.open.loc[idv, code] * 100
                        cost = np.maximum(price_amount*self.cost_ratio, 1)
                        if self.money >= (price_amount+cost): #资金满足开仓要求
                            buy_dict[code] += 1
                            price = self.open.loc[idv, code]
                            self.money -= (price * 100) 
                            self.cost_df.loc[idv, code] += cost  
                            self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                            print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = buy_dict[code]


    def pos_sub(self, idx, idv, code):
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        idv_pre = self.hold_pos.index[idx-1]
        hold_num = self.hold_pos.loc[idv_pre, code]
        if hold_num > 0: #持仓大于0
            self.sell_pos.loc[idv, code] = 0
            price = self.open.loc[idv, code].values 
            self.money += (price * 100 *hold_num)
                            
    def stop_ratio(self, idx, idv): #计算平均成本
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        for code in self.buy_pos.columns:
            if idx == 0:
                break
            id_init = idx - 1
            pos_i = self.buy_pos.loc[self.buy_pos.index[id_init],code]
            while (pos_i == 0) and (id_init > 0):               
                id_init = id_init - 1
                pos_i = self.buy_pos.loc[self.buy_pos.index[id_init],code]
            if pos_i == 0:
                continue
            id_initv = self.predict.index[id_init]
            price_init = self.open.loc[id_initv, code]
            price_now = self.close.loc[idv, code]
            ratio = 0
            if pos_i >= 1:
                ratio = price_now/price_init -1
            else:
                self.log_write('\nerror: 仓位为0，检查逻辑') 
                print('error: 仓位为0，检查逻辑') 
            #止盈平仓, #止损平仓  
            if (ratio >= self.stop_ret) or (ratio <= self.stop_loss):
                self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
                idv_pre = self.hold_pos.index[idx-1]
                hold_num = self.hold_pos.loc[idv_pre, code]
                if hold_num > 0: #持仓大于0
                    self.sell_pos.loc[idv, code] = hold_num
                    price = self.close.loc[idv, code]
                    self.money += (price * 100 *hold_num)
                    self.log_write(f'\n{idv}：{code}触发止损止盈平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                    print(f'{idv}：{code}触发止损止盈平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                    cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                    self.cost_df.loc[idv, code] += cost


class Pos_make_rank_thred_month_combine(Pos_make_rank_thred_step_combine):

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)

    def get_predict(self):
        combine_lst = [5,10,30]
        df = 0
        vol = 0
        for t in combine_lst:
            df_i = self.get_predict_t(t)
            df = df + df_i
            # df = df + df_i / df_i.rolling(5, min_periods=1).std()
            # vol = vol + (1/ df_i.rolling(5, min_periods=1).std())
        df = df.shift(1)
        df = df / len(combine_lst)
        # vol = vol.shift(1)
        df = df / vol
        return df

    def get_predict_t(self, t):    
        pred_list = []
        read_root = self.args.result_root[:-2]+str(t)
        file_list = sorted(os.listdir(read_root))
        for i in file_list:
            if i[-2:] == '01':
                df_i = pd.read_pickle(os.path.join(read_root, i, 'test_pred.pkl.gzip'))
                pred_list.append(df_i)
        df = pd.concat(pred_list, axis=0)   
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
         # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_rank_thred_month.__name__)
    #     return result_root


class Pos_make_series_thred_month_combine(Pos_make_rank_thred_month_combine):  

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root) 
        self.predict_vol = self.predict.rolling(self.vol_num, min_periods=1).std()   
        self.vol_num = 10
        self.open_thred = 2

    def pos_add(self, idx, idv):
        pred = self.predict.loc[idv] / (self.predict_vol.loc[idv]+0.0000001)
        pred = pred.sort_values(ascending=False)
        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处
                # if pred[code] >= self.open_thred: #满足开仓信号要求                       
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = buy_dict[code]
                
    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_series_thred_month.__name__)
    #     return result_root


class Pos_make_series_thred_step_combine(Pos_make_series_thred_month_combine):
    
    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)   

    def get_predict_t(self, t):
        read_root = self.args.result_root[:-2]+str(t)
        df = pd.read_pickle(os.path.join(read_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_series_thred_step.__name__)
    #     return result_root